{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s_qNSzzyaCbD"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jmjh290raIky"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J0Qjg6vuaHNt"
      },
      "source": [
        "# Transformer model for language understanding"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AOpGoE2T-YXS"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/text/tutorials/transformer\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003e\n",
        "    View on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/text/blob/master/docs/tutorials/transformer.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003e\n",
        "    Run in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/text/blob/master/docs/tutorials/transformer.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003e\n",
        "    View source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/text/docs/tutorials/transformer.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M-f8TnGpE_ex"
      },
      "source": [
        "This tutorial trains a \u003ca href=\"https://arxiv.org/abs/1706.03762\" class=\"external\"\u003eTransformer model\u003c/a\u003e to translate a [Portuguese to English dataset](https://www.tensorflow.org/datasets/catalog/ted_hrlr_translate#ted_hrlr_translatept_to_en). This is an advanced example that assumes knowledge of [text generation](https://www.tensorflow.org/text/tutorials/text_generation) and [attention](https://www.tensorflow.org/text/tutorials/nmt_with_attention).\n",
        "\n",
        "The core idea behind the Transformer model is *self-attention*—the ability to attend to different positions of the input sequence to compute a representation of that sequence. Transformer creates stacks of self-attention layers and is explained below in the sections *Scaled dot product attention* and *Multi-head attention*.\n",
        "\n",
        "A transformer model handles variable-sized input using stacks of self-attention layers instead of [RNNs](https://www.tensorflow.org/text/tutorials/text_classification_rnn) or [CNNs](https://www.tensorflow.org/tutorials/images/cnn). This general architecture has a number of advantages:\n",
        "\n",
        "* It makes no assumptions about the temporal/spatial relationships across the data. This is ideal for processing a set of objects (for example, [StarCraft units](https://deepmind.com/blog/alphastar-mastering-real-time-strategy-game-starcraft-ii/#block-8)).\n",
        "* Layer outputs can be calculated in parallel, instead of a series like an RNN.\n",
        "* Distant items can affect each other's output without passing through many RNN-steps, or convolution layers (see [Scene Memory Transformer](https://arxiv.org/pdf/1903.03878.pdf) for example).\n",
        "* It can learn long-range dependencies. This is a challenge in many sequence tasks.\n",
        "\n",
        "The downsides of this architecture are:\n",
        "\n",
        "* For a time-series, the output for a time-step is calculated from the *entire history* instead of only the inputs and current hidden-state. This _may_ be less efficient.   \n",
        "* If the input *does* have a  temporal/spatial relationship, like text, some positional encoding must be added or the model will effectively see a bag of words. \n",
        "\n",
        "After training the model in this notebook, you will be able to input a Portuguese sentence and return the English translation.\n",
        "\n",
        "\u003cimg src=\"https://www.tensorflow.org/images/tutorials/transformer/attention_map_portuguese.png\" width=\"800\" alt=\"Attention heatmap\"\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "swymtxpl7W7w"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XFG0NDRu5mYQ"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_datasets\n",
        "!pip install -U tensorflow-text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JjJJyJTZYebt"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "import logging\n",
        "import os\n",
        "import pathlib\n",
        "import re\n",
        "import string\n",
        "import sys\n",
        "import time\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow_text as text\n",
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pXzVhU34zWEU"
      },
      "outputs": [],
      "source": [
        "logging.getLogger('tensorflow').setLevel(logging.ERROR)  # suppress warnings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-cCvXbPkccV1"
      },
      "source": [
        "## Download the Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t4_Qt8W1hJE_"
      },
      "source": [
        "Use [TensorFlow datasets](https://www.tensorflow.org/datasets) to load the [Portuguese-English translation dataset](https://github.com/neulab/word-embeddings-for-nmt) from the [TED Talks Open Translation Project](https://www.ted.com/participate/translate).\n",
        "\n",
        "This dataset contains approximately 50000 training examples, 1100 validation examples, and 2000 test examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8q9t4FmN96eN"
      },
      "outputs": [],
      "source": [
        "examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en', with_info=True,\n",
        "                               as_supervised=True)\n",
        "train_examples, val_examples = examples['train'], examples['validation']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1OnwQZ2IpTbl"
      },
      "source": [
        "The `tf.data.Dataset` object returned by TensorFlow datasets yields pairs of text examples:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AtuMLTQGlHZ1"
      },
      "outputs": [],
      "source": [
        "for pt_examples, en_examples in train_examples.batch(3).take(1):\n",
        "  for pt in pt_examples.numpy():\n",
        "    print(pt.decode('utf-8'))\n",
        "\n",
        "  print()\n",
        "\n",
        "  for en in en_examples.numpy():\n",
        "    print(en.decode('utf-8'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eJxTd6aVnZyh"
      },
      "source": [
        "## Text tokenization \u0026 detokenization"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WlHAFNaopq6U"
      },
      "source": [
        "You can't train a model directly on text. The text needs to be converted to some numeric representation first. Typically, you convert the text to sequences of token IDs, which are used as indices into an embedding."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2sPvbCEePzOC"
      },
      "source": [
        "One popular implementation is demonstrated in the [Subword tokenizer tutorial](https://www.tensorflow.org/text/guide/subwords_tokenizer) builds subword tokenizers (`text.BertTokenizer`) optimized for this dataset and exports them in a [saved_model](https://www.tensorflow.org/guide/saved_model). \n",
        "\n",
        "Download and unzip and import the `saved_model`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QToMl0NanZPr"
      },
      "outputs": [],
      "source": [
        "model_name = \"ted_hrlr_translate_pt_en_converter\"\n",
        "tf.keras.utils.get_file(\n",
        "    f\"{model_name}.zip\",\n",
        "    f\"https://storage.googleapis.com/download.tensorflow.org/models/{model_name}.zip\",\n",
        "    cache_dir='.', cache_subdir='', extract=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h5dbGnPXnuI1"
      },
      "outputs": [],
      "source": [
        "tokenizers = tf.saved_model.load(model_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KQYc0CJTpnvg"
      },
      "source": [
        "The `tf.saved_model` contains two text tokenizers, one for English and one for Portuguese. Both have the same methods:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s-PCJijfcZ9_"
      },
      "outputs": [],
      "source": [
        "[item for item in dir(tokenizers.en) if not item.startswith('_')]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-pv_nc5lrAiD"
      },
      "source": [
        "The `tokenize` method converts a batch of strings to a padded-batch of token IDs. This method splits punctuation, lowercases and unicode-normalizes the input before tokenizing. That standardization is not visible here because the input data is already standardized."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0GV-smAAQJtM"
      },
      "outputs": [],
      "source": [
        "for en in en_examples.numpy():\n",
        "  print(en.decode('utf-8'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2fQJq1xB-tOn"
      },
      "outputs": [],
      "source": [
        "encoded = tokenizers.en.tokenize(en_examples)\n",
        "\n",
        "for row in encoded.to_list():\n",
        "  print(row)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dBRlikwDR2Lu"
      },
      "source": [
        "The `detokenize` method attempts to convert these token IDs back to human readable text: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tpqx1aug3W31"
      },
      "outputs": [],
      "source": [
        "round_trip = tokenizers.en.detokenize(encoded)\n",
        "for line in round_trip.numpy():\n",
        "  print(line.decode('utf-8'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5mGvytArL9g"
      },
      "source": [
        "The lower level `lookup` method converts from token-IDs to token text:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RpzRLzvIuN3R"
      },
      "outputs": [],
      "source": [
        "tokens = tokenizers.en.lookup(encoded)\n",
        "tokens"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5V8Ix_PNSfhV"
      },
      "source": [
        "Here you can see the \"subword\" aspect of the tokenizers. The word \"searchability\" is decomposed into \"search ##ability\" and the word \"serendipity\" into \"s ##ere ##nd ##ip ##ity\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fd1NWMxjfsDd"
      },
      "source": [
        "## Setup input pipeline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E7JHK7N7tNZy"
      },
      "source": [
        "To build an input pipeline suitable for training you'll apply some transformations to the dataset.\n",
        "\n",
        "This function will be used to encode the batches of raw text:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6shgzEck3FiV"
      },
      "outputs": [],
      "source": [
        "def tokenize_pairs(pt, en):\n",
        "    pt = tokenizers.pt.tokenize(pt)\n",
        "    # Convert from ragged to dense, padding with zeros.\n",
        "    pt = pt.to_tensor()\n",
        "\n",
        "    en = tokenizers.en.tokenize(en)\n",
        "    # Convert from ragged to dense, padding with zeros.\n",
        "    en = en.to_tensor()\n",
        "    return pt, en"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cf2HwC_0wG0v"
      },
      "source": [
        "Here's a simple input pipeline that processes, shuffles and batches the data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bcRp7VcQ5m6g"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = 20000\n",
        "BATCH_SIZE = 64"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BUN_jLBTwNxk"
      },
      "outputs": [],
      "source": [
        "def make_batches(ds):\n",
        "  return (\n",
        "      ds\n",
        "      .cache()\n",
        "      .shuffle(BUFFER_SIZE)\n",
        "      .batch(BATCH_SIZE)\n",
        "      .map(tokenize_pairs, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "      .prefetch(tf.data.AUTOTUNE))\n",
        "\n",
        "\n",
        "train_batches = make_batches(train_examples)\n",
        "val_batches = make_batches(val_examples)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nBQuibYA4n0n"
      },
      "source": [
        "## Positional encoding\n",
        "\n",
        "Attention layers see their input as a set of vectors, with no sequential order. This model also doesn't contain any recurrent or convolutional layers. Because of this a \"positional encoding\" is added to give the model some information about the relative position of the tokens in the sentence. \n",
        "\n",
        "The positional encoding vector is added to the embedding vector. Embeddings represent a token in a d-dimensional space where tokens with similar meaning will be closer to each other. But the embeddings do not encode the relative position of tokens in a sentence. So after adding the positional encoding, tokens will be closer to each other based on the *similarity of their meaning and their position in the sentence*, in the d-dimensional space.\n",
        "\n",
        "The formula for calculating the positional encoding is as follows:\n",
        "\n",
        "$$\\Large{PE_{(pos, 2i)} = \\sin(pos / 10000^{2i / d_{model}})} $$\n",
        "$$\\Large{PE_{(pos, 2i+1)} = \\cos(pos / 10000^{2i / d_{model}})} $$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WhIOZjMNKujn"
      },
      "outputs": [],
      "source": [
        "def get_angles(pos, i, d_model):\n",
        "  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
        "  return pos * angle_rates"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Rz82wEs5biZ"
      },
      "outputs": [],
      "source": [
        "def positional_encoding(position, d_model):\n",
        "  angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
        "                          np.arange(d_model)[np.newaxis, :],\n",
        "                          d_model)\n",
        "\n",
        "  # apply sin to even indices in the array; 2i\n",
        "  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
        "\n",
        "  # apply cos to odd indices in the array; 2i+1\n",
        "  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
        "\n",
        "  pos_encoding = angle_rads[np.newaxis, ...]\n",
        "\n",
        "  return tf.cast(pos_encoding, dtype=tf.float32)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1kLCla68EloE"
      },
      "outputs": [],
      "source": [
        "n, d = 2048, 512\n",
        "pos_encoding = positional_encoding(n, d)\n",
        "print(pos_encoding.shape)\n",
        "pos_encoding = pos_encoding[0]\n",
        "\n",
        "# Juggle the dimensions for the plot\n",
        "pos_encoding = tf.reshape(pos_encoding, (n, d//2, 2))\n",
        "pos_encoding = tf.transpose(pos_encoding, (2, 1, 0))\n",
        "pos_encoding = tf.reshape(pos_encoding, (d, n))\n",
        "\n",
        "plt.pcolormesh(pos_encoding, cmap='RdBu')\n",
        "plt.ylabel('Depth')\n",
        "plt.xlabel('Position')\n",
        "plt.colorbar()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_b4ou4TYqUN"
      },
      "source": [
        "## Masking"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s42Uydjkv0hF"
      },
      "source": [
        "Mask all the pad tokens in the batch of sequence. It ensures that the model does not treat padding as the input. The mask indicates where pad value `0` is present: it outputs a `1` at those locations, and a `0` otherwise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U2i8-e1s8ti9"
      },
      "outputs": [],
      "source": [
        "def create_padding_mask(seq):\n",
        "  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
        "\n",
        "  # add extra dimensions to add the padding\n",
        "  # to the attention logits.\n",
        "  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A7BYeBCNvi7n"
      },
      "outputs": [],
      "source": [
        "x = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])\n",
        "create_padding_mask(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z0hzukDBgVom"
      },
      "source": [
        "The look-ahead mask is used to mask the future tokens in a sequence. In other words, the mask indicates which entries should not be used.\n",
        "\n",
        "This means that to predict the third token, only the first and second token will be used. Similarly to predict the fourth token, only the first, second and the third tokens will be used and so on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dVxS8OPI9uI0"
      },
      "outputs": [],
      "source": [
        "def create_look_ahead_mask(size):\n",
        "  mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
        "  return mask  # (seq_len, seq_len)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yxKGuXxaBeeE"
      },
      "outputs": [],
      "source": [
        "x = tf.random.uniform((1, 3))\n",
        "temp = create_look_ahead_mask(x.shape[1])\n",
        "temp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xluDl5cXYy4y"
      },
      "source": [
        "## Scaled dot product attention"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vsxEE_-Wa1gF"
      },
      "source": [
        "\u003cimg src=\"https://www.tensorflow.org/images/tutorials/transformer/scaled_attention.png\" width=\"500\" alt=\"scaled_dot_product_attention\"\u003e\n",
        "\n",
        "The attention function used by the transformer takes three inputs: Q (query), K (key), V (value). The equation used to calculate the attention weights is:\n",
        "\n",
        "$$\\Large{Attention(Q, K, V) = softmax_k\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right) V} $$\n",
        "\n",
        "The dot-product attention is scaled by a factor of square root of the depth. This is done because for large values of depth, the dot product grows large in magnitude pushing the softmax function where it has small gradients resulting in a very hard softmax. \n",
        "\n",
        "For example, consider that `Q` and `K` have a mean of 0 and variance of 1. Their matrix multiplication will have a mean of 0 and variance of `dk`. So the *square root of `dk`* is used for scaling, so you get a consistent variance regardless of the value of `dk`. If the variance is too low the output may be too flat to optimize effectively. If the variance is too high the softmax may saturate at initialization making it difficult to learn.\n",
        "\n",
        "The mask is multiplied with -1e9 (close to negative infinity). This is done because the mask is summed with the scaled matrix multiplication of Q and K and is applied immediately before a softmax. The goal is to zero out these cells, and large negative inputs to softmax are near zero in the output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LazzUq3bJ5SH"
      },
      "outputs": [],
      "source": [
        "def scaled_dot_product_attention(q, k, v, mask):\n",
        "  \"\"\"Calculate the attention weights.\n",
        "  q, k, v must have matching leading dimensions.\n",
        "  k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.\n",
        "  The mask has different shapes depending on its type(padding or look ahead)\n",
        "  but it must be broadcastable for addition.\n",
        "\n",
        "  Args:\n",
        "    q: query shape == (..., seq_len_q, depth)\n",
        "    k: key shape == (..., seq_len_k, depth)\n",
        "    v: value shape == (..., seq_len_v, depth_v)\n",
        "    mask: Float tensor with shape broadcastable\n",
        "          to (..., seq_len_q, seq_len_k). Defaults to None.\n",
        "\n",
        "  Returns:\n",
        "    output, attention_weights\n",
        "  \"\"\"\n",
        "\n",
        "  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
        "\n",
        "  # scale matmul_qk\n",
        "  dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
        "  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
        "\n",
        "  # add the mask to the scaled tensor.\n",
        "  if mask is not None:\n",
        "    scaled_attention_logits += (mask * -1e9)\n",
        "\n",
        "  # softmax is normalized on the last axis (seq_len_k) so that the scores\n",
        "  # add up to 1.\n",
        "  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)\n",
        "\n",
        "  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
        "\n",
        "  return output, attention_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FiqETnhCkoXh"
      },
      "source": [
        "As the softmax normalization is done on K, its values decide the amount of importance given to Q.\n",
        "\n",
        "The output represents the multiplication of the attention weights and the V (value) vector. This ensures that the tokens you want to focus on are kept as-is and the irrelevant tokens are flushed out."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n90YjClyInFy"
      },
      "outputs": [],
      "source": [
        "def print_out(q, k, v):\n",
        "  temp_out, temp_attn = scaled_dot_product_attention(\n",
        "      q, k, v, None)\n",
        "  print('Attention weights are:')\n",
        "  print(temp_attn)\n",
        "  print('Output is:')\n",
        "  print(temp_out)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yAzUAf2DPlNt"
      },
      "outputs": [],
      "source": [
        "np.set_printoptions(suppress=True)\n",
        "\n",
        "temp_k = tf.constant([[10, 0, 0],\n",
        "                      [0, 10, 0],\n",
        "                      [0, 0, 10],\n",
        "                      [0, 0, 10]], dtype=tf.float32)  # (4, 3)\n",
        "\n",
        "temp_v = tf.constant([[1, 0],\n",
        "                      [10, 0],\n",
        "                      [100, 5],\n",
        "                      [1000, 6]], dtype=tf.float32)  # (4, 2)\n",
        "\n",
        "# This `query` aligns with the second `key`,\n",
        "# so the second `value` is returned.\n",
        "temp_q = tf.constant([[0, 10, 0]], dtype=tf.float32)  # (1, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zg6k-fGhgXra"
      },
      "outputs": [],
      "source": [
        "# This query aligns with a repeated key (third and fourth),\n",
        "# so all associated values get averaged.\n",
        "temp_q = tf.constant([[0, 0, 10]], dtype=tf.float32)  # (1, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UAq3YOzUgXhb"
      },
      "outputs": [],
      "source": [
        "# This query aligns equally with the first and second key,\n",
        "# so their values get averaged.\n",
        "temp_q = tf.constant([[10, 10, 0]], dtype=tf.float32)  # (1, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aOz-4_XIhaTP"
      },
      "source": [
        "Pass all the queries together."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6dlU8Tm-hYrF"
      },
      "outputs": [],
      "source": [
        "temp_q = tf.constant([[0, 0, 10],\n",
        "                      [0, 10, 0],\n",
        "                      [10, 10, 0]], dtype=tf.float32)  # (3, 3)\n",
        "print_out(temp_q, temp_k, temp_v)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kmzGPEy64qmA"
      },
      "source": [
        "## Multi-head attention"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fz5BMC8Kaoqo"
      },
      "source": [
        "\u003cimg src=\"https://www.tensorflow.org/images/tutorials/transformer/multi_head_attention.png\" width=\"500\" alt=\"multi-head attention\"\u003e\n",
        "\n",
        "\n",
        "Multi-head attention consists of four parts:\n",
        "*    Linear layers.\n",
        "*    Scaled dot-product attention.\n",
        "*    Final linear layer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPmbr6F1C-v_"
      },
      "source": [
        "Each multi-head attention block gets three inputs; Q (query), K (key), V (value). These are put through linear (Dense) layers before the multi-head attention function.\n",
        "\n",
        "In the diagram above `(K,Q,V)` are passed through sepearte linear (`Dense`) layers for each attention head. For simplicity/efficiency the code below implements this using a single dense layer with `num_heads` times as many outputs. The output is rearranged to a shape of `(batch, num_heads, ...)` before applying the attention function.\n",
        "\n",
        "The `scaled_dot_product_attention` function defined above is applied in a single call, broadcasted for efficiency. An appropriate mask must be used in the attention step.  The attention output for each head is then concatenated (using `tf.transpose`, and `tf.reshape`) and put through a final `Dense` layer.\n",
        "\n",
        "Instead of one single attention head, Q, K, and V are split into multiple heads because it allows the model to jointly attend to information from different representation subspaces at different positions. After the split each head has a reduced dimensionality, so the total computation cost is the same as a single head attention with full dimensionality."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BSV3PPKsYecw"
      },
      "outputs": [],
      "source": [
        "class MultiHeadAttention(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads):\n",
        "    super(MultiHeadAttention, self).__init__()\n",
        "    self.num_heads = num_heads\n",
        "    self.d_model = d_model\n",
        "\n",
        "    assert d_model % self.num_heads == 0\n",
        "\n",
        "    self.depth = d_model // self.num_heads\n",
        "\n",
        "    self.wq = tf.keras.layers.Dense(d_model)\n",
        "    self.wk = tf.keras.layers.Dense(d_model)\n",
        "    self.wv = tf.keras.layers.Dense(d_model)\n",
        "\n",
        "    self.dense = tf.keras.layers.Dense(d_model)\n",
        "\n",
        "  def split_heads(self, x, batch_size):\n",
        "    \"\"\"Split the last dimension into (num_heads, depth).\n",
        "    Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)\n",
        "    \"\"\"\n",
        "    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
        "    return tf.transpose(x, perm=[0, 2, 1, 3])\n",
        "\n",
        "  def call(self, v, k, q, mask):\n",
        "    batch_size = tf.shape(q)[0]\n",
        "\n",
        "    q = self.wq(q)  # (batch_size, seq_len, d_model)\n",
        "    k = self.wk(k)  # (batch_size, seq_len, d_model)\n",
        "    v = self.wv(v)  # (batch_size, seq_len, d_model)\n",
        "\n",
        "    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n",
        "    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n",
        "    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n",
        "\n",
        "    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)\n",
        "    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)\n",
        "    scaled_attention, attention_weights = scaled_dot_product_attention(\n",
        "        q, k, v, mask)\n",
        "\n",
        "    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)\n",
        "\n",
        "    concat_attention = tf.reshape(scaled_attention,\n",
        "                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)\n",
        "\n",
        "    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)\n",
        "\n",
        "    return output, attention_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0D8FJue5lDyZ"
      },
      "source": [
        "Create a `MultiHeadAttention` layer to try out. At each location in the sequence, `y`, the `MultiHeadAttention` runs all 8 attention heads across all other locations in the sequence, returning a new vector of the same length at each location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hu94p-_-2_BX"
      },
      "outputs": [],
      "source": [
        "temp_mha = MultiHeadAttention(d_model=512, num_heads=8)\n",
        "y = tf.random.uniform((1, 60, 512))  # (batch_size, encoder_sequence, d_model)\n",
        "out, attn = temp_mha(y, k=y, q=y, mask=None)\n",
        "out.shape, attn.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RdDqGayx67vv"
      },
      "source": [
        "## Point wise feed forward network"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gBqzJXGfHK3X"
      },
      "source": [
        "Point wise feed forward network consists of two fully-connected layers with a ReLU activation in between."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ET7xLt0yCT6Z"
      },
      "outputs": [],
      "source": [
        "def point_wise_feed_forward_network(d_model, dff):\n",
        "  return tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)\n",
        "      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
        "  ])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mytb1lPyOHLB"
      },
      "outputs": [],
      "source": [
        "sample_ffn = point_wise_feed_forward_network(512, 2048)\n",
        "sample_ffn(tf.random.uniform((64, 50, 512))).shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7e7hKcxn6-zd"
      },
      "source": [
        "## Encoder and decoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yScbC0MUH8dS"
      },
      "source": [
        "\u003cimg src=\"https://www.tensorflow.org/images/tutorials/transformer/transformer.png\" width=\"600\" alt=\"transformer\"\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfYJG-Kvgwy2"
      },
      "source": [
        "The transformer model follows the same general pattern as a standard [sequence to sequence with attention model](https://www.tensorflow.org/text/tutorials/nmt_with_attention.ipynb). \n",
        "\n",
        "* The input sentence is passed through `N` encoder layers that generates an output for each token in the sequence.\n",
        "* The decoder attends to the encoder's output and its own input (self-attention) to predict the next word. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QFv-FNYUmvpn"
      },
      "source": [
        "### Encoder layer\n",
        "\n",
        "Each encoder layer consists of sublayers:\n",
        "\n",
        "1.   Multi-head attention (with padding mask) \n",
        "2.    Point wise feed forward networks. \n",
        "\n",
        "Each of these sublayers has a residual connection around it followed by a layer normalization. Residual connections help in avoiding the vanishing gradient problem in deep networks.\n",
        "\n",
        "The output of each sublayer is `LayerNorm(x + Sublayer(x))`. The normalization is done on the `d_model` (last) axis. There are N encoder layers in the transformer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ncyS-Ms3i2x_"
      },
      "outputs": [],
      "source": [
        "class EncoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
        "    super(EncoderLayer, self).__init__()\n",
        "\n",
        "    self.mha = MultiHeadAttention(d_model, num_heads)\n",
        "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "\n",
        "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, training, mask):\n",
        "\n",
        "    attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)\n",
        "    attn_output = self.dropout1(attn_output, training=training)\n",
        "    out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)\n",
        "\n",
        "    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
        "    ffn_output = self.dropout2(ffn_output, training=training)\n",
        "    out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)\n",
        "\n",
        "    return out2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AzZRXdO0mI48"
      },
      "outputs": [],
      "source": [
        "sample_encoder_layer = EncoderLayer(512, 8, 2048)\n",
        "\n",
        "sample_encoder_layer_output = sample_encoder_layer(\n",
        "    tf.random.uniform((64, 43, 512)), False, None)\n",
        "\n",
        "sample_encoder_layer_output.shape  # (batch_size, input_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6LO_48Owmx_o"
      },
      "source": [
        "### Decoder layer\n",
        "\n",
        "Each decoder layer consists of sublayers:\n",
        "\n",
        "1.   Masked multi-head attention (with look ahead mask and padding mask)\n",
        "2.   Multi-head attention (with padding mask). V (value) and K (key) receive the *encoder output* as inputs. Q (query) receives the *output from the masked multi-head attention sublayer.*\n",
        "3.   Point wise feed forward networks\n",
        "\n",
        "Each of these sublayers has a residual connection around it followed by a layer normalization. The output of each sublayer is `LayerNorm(x + Sublayer(x))`. The normalization is done on the `d_model` (last) axis.\n",
        "\n",
        "There are N decoder layers in the transformer.\n",
        "\n",
        "As Q receives the output from decoder's first attention block, and K receives the encoder output, the attention weights represent the importance given to the decoder's input based on the encoder's output. In other words, the decoder predicts the next token by looking at the encoder output and self-attending to its own output. See the demonstration above in the scaled dot product attention section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9SoX0-vd1hue"
      },
      "outputs": [],
      "source": [
        "class DecoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, dff, rate=0.1):\n",
        "    super(DecoderLayer, self).__init__()\n",
        "\n",
        "    self.mha1 = MultiHeadAttention(d_model, num_heads)\n",
        "    self.mha2 = MultiHeadAttention(d_model, num_heads)\n",
        "\n",
        "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "\n",
        "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout3 = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, enc_output, training,\n",
        "           look_ahead_mask, padding_mask):\n",
        "    # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
        "\n",
        "    attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)\n",
        "    attn1 = self.dropout1(attn1, training=training)\n",
        "    out1 = self.layernorm1(attn1 + x)\n",
        "\n",
        "    attn2, attn_weights_block2 = self.mha2(\n",
        "        enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)\n",
        "    attn2 = self.dropout2(attn2, training=training)\n",
        "    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)\n",
        "\n",
        "    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n",
        "    ffn_output = self.dropout3(ffn_output, training=training)\n",
        "    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n",
        "\n",
        "    return out3, attn_weights_block1, attn_weights_block2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ne2Bqx8k71l0"
      },
      "outputs": [],
      "source": [
        "sample_decoder_layer = DecoderLayer(512, 8, 2048)\n",
        "\n",
        "sample_decoder_layer_output, _, _ = sample_decoder_layer(\n",
        "    tf.random.uniform((64, 50, 512)), sample_encoder_layer_output,\n",
        "    False, None, None)\n",
        "\n",
        "sample_decoder_layer_output.shape  # (batch_size, target_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SE1H51Ajm0q1"
      },
      "source": [
        "### Encoder\n",
        "\n",
        "The `Encoder` consists of:\n",
        "1.   Input Embedding\n",
        "2.   Positional Encoding\n",
        "3.   N encoder layers\n",
        "\n",
        "The input is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the encoder layers. The output of the encoder is the input to the decoder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpEox7gJ8FCI"
      },
      "outputs": [],
      "source": [
        "class Encoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,\n",
        "               maximum_position_encoding, rate=0.1):\n",
        "    super(Encoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "\n",
        "    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)\n",
        "    self.pos_encoding = positional_encoding(maximum_position_encoding,\n",
        "                                            self.d_model)\n",
        "\n",
        "    self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate)\n",
        "                       for _ in range(num_layers)]\n",
        "\n",
        "    self.dropout = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, training, mask):\n",
        "\n",
        "    seq_len = tf.shape(x)[1]\n",
        "\n",
        "    # adding embedding and position encoding.\n",
        "    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)\n",
        "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "    x += self.pos_encoding[:, :seq_len, :]\n",
        "\n",
        "    x = self.dropout(x, training=training)\n",
        "\n",
        "    for i in range(self.num_layers):\n",
        "      x = self.enc_layers[i](x, training, mask)\n",
        "\n",
        "    return x  # (batch_size, input_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8QG9nueFQKXx"
      },
      "outputs": [],
      "source": [
        "sample_encoder = Encoder(num_layers=2, d_model=512, num_heads=8,\n",
        "                         dff=2048, input_vocab_size=8500,\n",
        "                         maximum_position_encoding=10000)\n",
        "temp_input = tf.random.uniform((64, 62), dtype=tf.int64, minval=0, maxval=200)\n",
        "\n",
        "sample_encoder_output = sample_encoder(temp_input, training=False, mask=None)\n",
        "\n",
        "print(sample_encoder_output.shape)  # (batch_size, input_seq_len, d_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p-uO6ls8m2O5"
      },
      "source": [
        "### Decoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZtT7PKzrXkNr"
      },
      "source": [
        " The `Decoder` consists of:\n",
        "1.   Output Embedding\n",
        "2.   Positional Encoding\n",
        "3.   N decoder layers\n",
        "\n",
        "The target is put through an embedding which is summed with the positional encoding. The output of this summation is the input to the decoder layers. The output of the decoder is the input to the final linear layer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5_d5-PLQXwY"
      },
      "outputs": [],
      "source": [
        "class Decoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,\n",
        "               maximum_position_encoding, rate=0.1):\n",
        "    super(Decoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "\n",
        "    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n",
        "    self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)\n",
        "\n",
        "    self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate)\n",
        "                       for _ in range(num_layers)]\n",
        "    self.dropout = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, enc_output, training,\n",
        "           look_ahead_mask, padding_mask):\n",
        "\n",
        "    seq_len = tf.shape(x)[1]\n",
        "    attention_weights = {}\n",
        "\n",
        "    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n",
        "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "    x += self.pos_encoding[:, :seq_len, :]\n",
        "\n",
        "    x = self.dropout(x, training=training)\n",
        "\n",
        "    for i in range(self.num_layers):\n",
        "      x, block1, block2 = self.dec_layers[i](x, enc_output, training,\n",
        "                                             look_ahead_mask, padding_mask)\n",
        "\n",
        "      attention_weights[f'decoder_layer{i+1}_block1'] = block1\n",
        "      attention_weights[f'decoder_layer{i+1}_block2'] = block2\n",
        "\n",
        "    # x.shape == (batch_size, target_seq_len, d_model)\n",
        "    return x, attention_weights"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a1jXoAMRZyvu"
      },
      "outputs": [],
      "source": [
        "sample_decoder = Decoder(num_layers=2, d_model=512, num_heads=8,\n",
        "                         dff=2048, target_vocab_size=8000,\n",
        "                         maximum_position_encoding=5000)\n",
        "temp_input = tf.random.uniform((64, 26), dtype=tf.int64, minval=0, maxval=200)\n",
        "\n",
        "output, attn = sample_decoder(temp_input,\n",
        "                              enc_output=sample_encoder_output,\n",
        "                              training=False,\n",
        "                              look_ahead_mask=None,\n",
        "                              padding_mask=None)\n",
        "\n",
        "output.shape, attn['decoder_layer2_block2'].shape"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y54xnJnuYgJ7"
      },
      "source": [
        "## Create the Transformer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uERO1y54cOKq"
      },
      "source": [
        "Transformer consists of the encoder, decoder and a final linear layer. The output of the decoder is the input to the linear layer and its output is returned."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PED3bIpOYkBu"
      },
      "outputs": [],
      "source": [
        "class Transformer(tf.keras.Model):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,\n",
        "               target_vocab_size, pe_input, pe_target, rate=0.1):\n",
        "    super().__init__()\n",
        "    self.encoder = Encoder(num_layers, d_model, num_heads, dff,\n",
        "                             input_vocab_size, pe_input, rate)\n",
        "\n",
        "    self.decoder = Decoder(num_layers, d_model, num_heads, dff,\n",
        "                           target_vocab_size, pe_target, rate)\n",
        "\n",
        "    self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
        "\n",
        "  def call(self, inputs, training):\n",
        "    # Keras models prefer if you pass all your inputs in the first argument\n",
        "    inp, tar = inputs\n",
        "\n",
        "    enc_padding_mask, look_ahead_mask, dec_padding_mask = self.create_masks(inp, tar)\n",
        "\n",
        "    enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)\n",
        "\n",
        "    # dec_output.shape == (batch_size, tar_seq_len, d_model)\n",
        "    dec_output, attention_weights = self.decoder(\n",
        "        tar, enc_output, training, look_ahead_mask, dec_padding_mask)\n",
        "\n",
        "    final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)\n",
        "\n",
        "    return final_output, attention_weights\n",
        "\n",
        "  def create_masks(self, inp, tar):\n",
        "    # Encoder padding mask\n",
        "    enc_padding_mask = create_padding_mask(inp)\n",
        "\n",
        "    # Used in the 2nd attention block in the decoder.\n",
        "    # This padding mask is used to mask the encoder outputs.\n",
        "    dec_padding_mask = create_padding_mask(inp)\n",
        "\n",
        "    # Used in the 1st attention block in the decoder.\n",
        "    # It is used to pad and mask future tokens in the input received by\n",
        "    # the decoder.\n",
        "    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n",
        "    dec_target_padding_mask = create_padding_mask(tar)\n",
        "    look_ahead_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
        "\n",
        "    return enc_padding_mask, look_ahead_mask, dec_padding_mask"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tJ4fbQcIkHW1"
      },
      "outputs": [],
      "source": [
        "sample_transformer = Transformer(\n",
        "    num_layers=2, d_model=512, num_heads=8, dff=2048,\n",
        "    input_vocab_size=8500, target_vocab_size=8000,\n",
        "    pe_input=10000, pe_target=6000)\n",
        "\n",
        "temp_input = tf.random.uniform((64, 38), dtype=tf.int64, minval=0, maxval=200)\n",
        "temp_target = tf.random.uniform((64, 36), dtype=tf.int64, minval=0, maxval=200)\n",
        "\n",
        "fn_out, _ = sample_transformer([temp_input, temp_target], training=False)\n",
        "\n",
        "fn_out.shape  # (batch_size, tar_seq_len, target_vocab_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wsINyf1VEQLC"
      },
      "source": [
        "## Set hyperparameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zVjWCxFNcgbt"
      },
      "source": [
        "To keep this example small and relatively fast, the values for `num_layers, d_model, dff` have been reduced. \n",
        "\n",
        "The base model described in the [paper](https://arxiv.org/abs/1706.03762) used: `num_layers=6, d_model=512, dff=2048`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnJn5SLA2ahP"
      },
      "outputs": [],
      "source": [
        "num_layers = 4\n",
        "d_model = 128\n",
        "dff = 512\n",
        "num_heads = 8\n",
        "dropout_rate = 0.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xYEGhEOtzn5W"
      },
      "source": [
        "## Optimizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GOmWW--yP3zx"
      },
      "source": [
        "Use the Adam optimizer with a custom learning rate scheduler according to the formula in the [paper](https://arxiv.org/abs/1706.03762).\n",
        "\n",
        "$$\\Large{lrate = d_{model}^{-0.5} * \\min(step{\\_}num^{-0.5}, step{\\_}num \\cdot warmup{\\_}steps^{-1.5})}$$\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iYQdOO1axwEI"
      },
      "outputs": [],
      "source": [
        "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
        "  def __init__(self, d_model, warmup_steps=4000):\n",
        "    super(CustomSchedule, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.d_model = tf.cast(self.d_model, tf.float32)\n",
        "\n",
        "    self.warmup_steps = warmup_steps\n",
        "\n",
        "  def __call__(self, step):\n",
        "    arg1 = tf.math.rsqrt(step)\n",
        "    arg2 = step * (self.warmup_steps ** -1.5)\n",
        "\n",
        "    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7r4scdulztRx"
      },
      "outputs": [],
      "source": [
        "learning_rate = CustomSchedule(d_model)\n",
        "\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,\n",
        "                                     epsilon=1e-9)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f33ZCgvHpPdG"
      },
      "outputs": [],
      "source": [
        "temp_learning_rate_schedule = CustomSchedule(d_model)\n",
        "\n",
        "plt.plot(temp_learning_rate_schedule(tf.range(40000, dtype=tf.float32)))\n",
        "plt.ylabel(\"Learning Rate\")\n",
        "plt.xlabel(\"Train Step\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YgkDE7hzo8r5"
      },
      "source": [
        "## Loss and metrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oxGJtoDuYIHL"
      },
      "source": [
        "Since the target sequences are padded, it is important to apply a padding mask when calculating the loss."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MlhsJMm0TW_B"
      },
      "outputs": [],
      "source": [
        "loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=True, reduction='none')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "67oqVHiT0Eiu"
      },
      "outputs": [],
      "source": [
        "def loss_function(real, pred):\n",
        "  mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
        "  loss_ = loss_object(real, pred)\n",
        "\n",
        "  mask = tf.cast(mask, dtype=loss_.dtype)\n",
        "  loss_ *= mask\n",
        "\n",
        "  return tf.reduce_sum(loss_)/tf.reduce_sum(mask)\n",
        "\n",
        "\n",
        "def accuracy_function(real, pred):\n",
        "  accuracies = tf.equal(real, tf.argmax(pred, axis=2))\n",
        "\n",
        "  mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
        "  accuracies = tf.math.logical_and(mask, accuracies)\n",
        "\n",
        "  accuracies = tf.cast(accuracies, dtype=tf.float32)\n",
        "  mask = tf.cast(mask, dtype=tf.float32)\n",
        "  return tf.reduce_sum(accuracies)/tf.reduce_sum(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "phlyxMnm-Tpx"
      },
      "outputs": [],
      "source": [
        "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
        "train_accuracy = tf.keras.metrics.Mean(name='train_accuracy')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aeHumfr7zmMa"
      },
      "source": [
        "## Training and checkpointing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UiysUa--4tOU"
      },
      "outputs": [],
      "source": [
        "transformer = Transformer(\n",
        "    num_layers=num_layers,\n",
        "    d_model=d_model,\n",
        "    num_heads=num_heads,\n",
        "    dff=dff,\n",
        "    input_vocab_size=tokenizers.pt.get_vocab_size().numpy(),\n",
        "    target_vocab_size=tokenizers.en.get_vocab_size().numpy(),\n",
        "    pe_input=1000,\n",
        "    pe_target=1000,\n",
        "    rate=dropout_rate)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fzuf06YZp66w"
      },
      "source": [
        "Create the checkpoint path and the checkpoint manager. This will be used to save checkpoints every `n` epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hNhuYfllndLZ"
      },
      "outputs": [],
      "source": [
        "checkpoint_path = \"./checkpoints/train\"\n",
        "\n",
        "ckpt = tf.train.Checkpoint(transformer=transformer,\n",
        "                           optimizer=optimizer)\n",
        "\n",
        "ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
        "\n",
        "# if a checkpoint exists, restore the latest checkpoint.\n",
        "if ckpt_manager.latest_checkpoint:\n",
        "  ckpt.restore(ckpt_manager.latest_checkpoint)\n",
        "  print('Latest checkpoint restored!!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Di_Yaa1gf9r"
      },
      "source": [
        "The target is divided into tar_inp and tar_real. tar_inp is passed as an input to the decoder. `tar_real` is that same input shifted by 1: At each location in `tar_input`, `tar_real` contains the  next token that should be predicted.\n",
        "\n",
        "For example, `sentence` = \"SOS A lion in the jungle is sleeping EOS\"\n",
        "\n",
        "`tar_inp` =  \"SOS A lion in the jungle is sleeping\"\n",
        "\n",
        "`tar_real` = \"A lion in the jungle is sleeping EOS\"\n",
        "\n",
        "The transformer is an auto-regressive model: it makes predictions one part at a time, and uses its output so far to decide what to do next. \n",
        "\n",
        "During training this example uses teacher-forcing (like in the [text generation tutorial](https://www.tensorflow.org/text/tutorials/text_generation)). Teacher forcing is passing the true output to the next time step regardless of what the model predicts at the current time step.\n",
        "\n",
        "As the transformer predicts each token, *self-attention* allows it to look at the previous tokens in the input sequence to better predict the next token.\n",
        "\n",
        "To prevent the model from peeking at the expected output the model uses a look-ahead mask."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LKpoA6q1sJFj"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 20"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iJwmp9OE29oj"
      },
      "outputs": [],
      "source": [
        "# The @tf.function trace-compiles train_step into a TF graph for faster\n",
        "# execution. The function specializes to the precise shape of the argument\n",
        "# tensors. To avoid re-tracing due to the variable sequence lengths or variable\n",
        "# batch sizes (the last batch is smaller), use input_signature to specify\n",
        "# more generic shapes.\n",
        "\n",
        "train_step_signature = [\n",
        "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
        "    tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
        "]\n",
        "\n",
        "\n",
        "@tf.function(input_signature=train_step_signature)\n",
        "def train_step(inp, tar):\n",
        "  tar_inp = tar[:, :-1]\n",
        "  tar_real = tar[:, 1:]\n",
        "\n",
        "  with tf.GradientTape() as tape:\n",
        "    predictions, _ = transformer([inp, tar_inp],\n",
        "                                 training = True)\n",
        "    loss = loss_function(tar_real, predictions)\n",
        "\n",
        "  gradients = tape.gradient(loss, transformer.trainable_variables)\n",
        "  optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
        "\n",
        "  train_loss(loss)\n",
        "  train_accuracy(accuracy_function(tar_real, predictions))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qM2PDWGDJ_8V"
      },
      "source": [
        "Portuguese is used as the input language and English is the target language."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bbvmaKNiznHZ"
      },
      "outputs": [],
      "source": [
        "for epoch in range(EPOCHS):\n",
        "  start = time.time()\n",
        "\n",
        "  train_loss.reset_states()\n",
        "  train_accuracy.reset_states()\n",
        "\n",
        "  # inp -\u003e portuguese, tar -\u003e english\n",
        "  for (batch, (inp, tar)) in enumerate(train_batches):\n",
        "    train_step(inp, tar)\n",
        "\n",
        "    if batch % 50 == 0:\n",
        "      print(f'Epoch {epoch + 1} Batch {batch} Loss {train_loss.result():.4f} Accuracy {train_accuracy.result():.4f}')\n",
        "\n",
        "  if (epoch + 1) % 5 == 0:\n",
        "    ckpt_save_path = ckpt_manager.save()\n",
        "    print(f'Saving checkpoint for epoch {epoch+1} at {ckpt_save_path}')\n",
        "\n",
        "  print(f'Epoch {epoch + 1} Loss {train_loss.result():.4f} Accuracy {train_accuracy.result():.4f}')\n",
        "\n",
        "  print(f'Time taken for 1 epoch: {time.time() - start:.2f} secs\\n')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QfcsSWswSdGV"
      },
      "source": [
        "### Run inference"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6APsFrgImLW"
      },
      "source": [
        "The following steps are used for inference:\n",
        "\n",
        "* Encode the input sentence using the Portuguese tokenizer (`tokenizers.pt`). This is the encoder input.\n",
        "* The decoder input is initialized to the `[START]` token.\n",
        "* Calculate the padding masks and the look ahead masks.\n",
        "* The `decoder` then outputs the predictions by looking at the `encoder output` and its own output (self-attention).\n",
        "* Concatenate the predicted token to the decoder input and pass it to the decoder.\n",
        "* In this approach, the decoder predicts the next token based on the previous tokens it predicted."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-FQmQwtv9-kk"
      },
      "source": [
        "Note: The model is optimized for _efficient training_ and makes a next-token prediction for each token in the output simultaneously. This is redundant during inference, and only the last prediction is used.  This model can be made more efficient for inference if you only calculate the last prediction when running in inference mode (`training=False`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5buvMlnvyrFm"
      },
      "outputs": [],
      "source": [
        "class Translator(tf.Module):\n",
        "  def __init__(self, tokenizers, transformer):\n",
        "    self.tokenizers = tokenizers\n",
        "    self.transformer = transformer\n",
        "\n",
        "  def __call__(self, sentence, max_length=20):\n",
        "    # input sentence is portuguese, hence adding the start and end token\n",
        "    assert isinstance(sentence, tf.Tensor)\n",
        "    if len(sentence.shape) == 0:\n",
        "      sentence = sentence[tf.newaxis]\n",
        "\n",
        "    sentence = self.tokenizers.pt.tokenize(sentence).to_tensor()\n",
        "    \n",
        "    encoder_input = sentence\n",
        "\n",
        "    # as the target is english, the first token to the transformer should be the\n",
        "    # english start token.\n",
        "    start_end = self.tokenizers.en.tokenize([''])[0]\n",
        "    start = start_end[0][tf.newaxis]\n",
        "    end = start_end[1][tf.newaxis]\n",
        "\n",
        "    # `tf.TensorArray` is required here (instead of a python list) so that the\n",
        "    # dynamic-loop can be traced by `tf.function`.\n",
        "    output_array = tf.TensorArray(dtype=tf.int64, size=0, dynamic_size=True)\n",
        "    output_array = output_array.write(0, start)\n",
        "    \n",
        "    for i in tf.range(max_length):\n",
        "      output = tf.transpose(output_array.stack())\n",
        "      predictions, _ = self.transformer([encoder_input, output], training=False)\n",
        "      \n",
        "      # select the last token from the seq_len dimension\n",
        "      predictions = predictions[:, -1:, :]  # (batch_size, 1, vocab_size)\n",
        "\n",
        "      predicted_id = tf.argmax(predictions, axis=-1)\n",
        "\n",
        "      # concatentate the predicted_id to the output which is given to the decoder\n",
        "      # as its input.\n",
        "      output_array = output_array.write(i+1, predicted_id[0])\n",
        "\n",
        "      if predicted_id == end:\n",
        "        break\n",
        "\n",
        "    output = tf.transpose(output_array.stack())\n",
        "    # output.shape (1, tokens)\n",
        "    text = tokenizers.en.detokenize(output)[0]  # shape: ()\n",
        "\n",
        "    tokens = tokenizers.en.lookup(output)[0]\n",
        "\n",
        "    # `tf.function` prevents us from using the attention_weights that were\n",
        "    # calculated on the last iteration of the loop. So recalculate them outside\n",
        "    # the loop.\n",
        "    _, attention_weights = self.transformer([encoder_input, output[:,:-1]], training=False)\n",
        "\n",
        "    return text, tokens, attention_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ofUWszmY3szZ"
      },
      "source": [
        "Create an instance of this `Translator` class, and try it out a few times:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4OR2D4EXeIRY"
      },
      "outputs": [],
      "source": [
        "translator = Translator(tokenizers, transformer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lU2_yG_vBGza"
      },
      "outputs": [],
      "source": [
        "def print_translation(sentence, tokens, ground_truth):\n",
        "  print(f'{\"Input:\":15s}: {sentence}')\n",
        "  print(f'{\"Prediction\":15s}: {tokens.numpy().decode(\"utf-8\")}')\n",
        "  print(f'{\"Ground truth\":15s}: {ground_truth}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YsxrAlvFG8SZ"
      },
      "outputs": [],
      "source": [
        "sentence = \"este é um problema que temos que resolver.\"\n",
        "ground_truth = \"this is a problem we have to solve .\"\n",
        "\n",
        "translated_text, translated_tokens, attention_weights = translator(\n",
        "    tf.constant(sentence))\n",
        "print_translation(sentence, translated_text, ground_truth)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7EH5y_aqI4t1"
      },
      "outputs": [],
      "source": [
        "sentence = \"os meus vizinhos ouviram sobre esta ideia.\"\n",
        "ground_truth = \"and my neighboring homes heard about this idea .\"\n",
        "\n",
        "translated_text, translated_tokens, attention_weights = translator(\n",
        "    tf.constant(sentence))\n",
        "print_translation(sentence, translated_text, ground_truth)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J-hVCTSUMlkb"
      },
      "outputs": [],
      "source": [
        "sentence = \"vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram.\"\n",
        "ground_truth = \"so i \\'ll just share with you some stories very quickly of some magical things that have happened .\"\n",
        "\n",
        "translated_text, translated_tokens, attention_weights = translator(\n",
        "    tf.constant(sentence))\n",
        "print_translation(sentence, translated_text, ground_truth)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3EQiFUC--Ds"
      },
      "source": [
        "## Attention plots"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hHV2pdXHGz-0"
      },
      "source": [
        "The `Translator` class returns a dictionary of attention maps you can use to visualize the internal working of the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t-kFyiOLH0xg"
      },
      "outputs": [],
      "source": [
        "sentence = \"este é o primeiro livro que eu fiz.\"\n",
        "ground_truth = \"this is the first book i've ever done.\"\n",
        "\n",
        "translated_text, translated_tokens, attention_weights = translator(\n",
        "    tf.constant(sentence))\n",
        "print_translation(sentence, translated_text, ground_truth)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CcI4DxAK5EHY"
      },
      "outputs": [],
      "source": [
        "def plot_attention_head(in_tokens, translated_tokens, attention):\n",
        "  # The plot is of the attention when a token was generated.\n",
        "  # The model didn't generate `\u003cSTART\u003e` in the output. Skip it.\n",
        "  translated_tokens = translated_tokens[1:]\n",
        "\n",
        "  ax = plt.gca()\n",
        "  ax.matshow(attention)\n",
        "  ax.set_xticks(range(len(in_tokens)))\n",
        "  ax.set_yticks(range(len(translated_tokens)))\n",
        "\n",
        "  labels = [label.decode('utf-8') for label in in_tokens.numpy()]\n",
        "  ax.set_xticklabels(\n",
        "      labels, rotation=90)\n",
        "\n",
        "  labels = [label.decode('utf-8') for label in translated_tokens.numpy()]\n",
        "  ax.set_yticklabels(labels)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_KY4c2cryuxY"
      },
      "outputs": [],
      "source": [
        "head = 0\n",
        "# shape: (batch=1, num_heads, seq_len_q, seq_len_k)\n",
        "attention_heads = tf.squeeze(\n",
        "  attention_weights['decoder_layer4_block2'], 0)\n",
        "attention = attention_heads[head]\n",
        "attention.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XdxmakWE6Om3"
      },
      "outputs": [],
      "source": [
        "in_tokens = tf.convert_to_tensor([sentence])\n",
        "in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()\n",
        "in_tokens = tokenizers.pt.lookup(in_tokens)[0]\n",
        "in_tokens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hVdPSfecmrpj"
      },
      "outputs": [],
      "source": [
        "translated_tokens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtzyKCFamm4N"
      },
      "outputs": [],
      "source": [
        "plot_attention_head(in_tokens, translated_tokens, attention)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MBliB-PCzNK3"
      },
      "outputs": [],
      "source": [
        "def plot_attention_weights(sentence, translated_tokens, attention_heads):\n",
        "  in_tokens = tf.convert_to_tensor([sentence])\n",
        "  in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()\n",
        "  in_tokens = tokenizers.pt.lookup(in_tokens)[0]\n",
        "  in_tokens\n",
        "\n",
        "  fig = plt.figure(figsize=(16, 8))\n",
        "\n",
        "  for h, head in enumerate(attention_heads):\n",
        "    ax = fig.add_subplot(2, 4, h+1)\n",
        "\n",
        "    plot_attention_head(in_tokens, translated_tokens, head)\n",
        "\n",
        "    ax.set_xlabel(f'Head {h+1}')\n",
        "\n",
        "  plt.tight_layout()\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pyRQi7944wru"
      },
      "outputs": [],
      "source": [
        "plot_attention_weights(sentence, translated_tokens,\n",
        "                       attention_weights['decoder_layer4_block2'][0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MZJirKUtikTt"
      },
      "source": [
        "The model does okay on unfamiliar words. Neither \"triceratops\" or \"encyclopedia\" are in the input dataset and the model almost learns to transliterate them, even without a shared vocabulary:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9cxysY7uh3jg"
      },
      "outputs": [],
      "source": [
        "sentence = \"Eu li sobre triceratops na enciclopédia.\"\n",
        "ground_truth = \"I read about triceratops in the encyclopedia.\"\n",
        "\n",
        "translated_text, translated_tokens, attention_weights = translator(\n",
        "    tf.constant(sentence))\n",
        "print_translation(sentence, translated_text, ground_truth)\n",
        "\n",
        "plot_attention_weights(sentence, translated_tokens,\n",
        "                       attention_weights['decoder_layer4_block2'][0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mOyiOetL2l60"
      },
      "source": [
        "## Export"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YTK3g2UL2oMc"
      },
      "source": [
        "That inference model is working, so next you'll export it as a `tf.saved_model`.\n",
        "\n",
        "To do that, wrap it in yet another `tf.Module` sub-class, this time with a `tf.function` on the `__call__` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GRmzkibLusQi"
      },
      "outputs": [],
      "source": [
        "class ExportTranslator(tf.Module):\n",
        "  def __init__(self, translator):\n",
        "    self.translator = translator\n",
        "\n",
        "  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])\n",
        "  def __call__(self, sentence):\n",
        "    (result, \n",
        "     tokens,\n",
        "     attention_weights) = self.translator(sentence, max_length=100)\n",
        "    \n",
        "    return result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O9f_pmEA4kql"
      },
      "source": [
        "In the above `tf.function` only the output sentence is returned. Thanks to the [non-strict execution](https://tensorflow.org/guide/intro_to_graphs) in `tf.function` any unnecessary values are never computed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EfomoJDP2n5n"
      },
      "outputs": [],
      "source": [
        "translator = ExportTranslator(translator)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SUfoCWPS9LuB"
      },
      "source": [
        "Since the model is decoding the predictions using `tf.argmax` the predictions are deterministic. The original model and one reloaded from its `SavedModel` should give identical predictions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hAlqyycz3IYL"
      },
      "outputs": [],
      "source": [
        "translator(\"este é o primeiro livro que eu fiz.\").numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ar3LO-Vuvlnv"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(translator, export_dir='translator')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8WUflwyT1SEF"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load('translator')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-sBTBWwR1XMr"
      },
      "outputs": [],
      "source": [
        "reloaded(\"este é o primeiro livro que eu fiz.\").numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RqQ1fIsLwkGE"
      },
      "source": [
        "## Summary\n",
        "\n",
        "In this tutorial, you learned about positional encoding, multi-head attention, the importance of masking and how to create a transformer.\n",
        "\n",
        "Try using a different dataset to train the transformer. You can also create the base transformer or transformer XL by changing the hyperparameters above. You can also use the layers defined here to create [BERT](https://arxiv.org/abs/1810.04805) and train state of the art models. Furthermore, you can implement beam search to get better predictions."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "transformer.ipynb",
      "provenance": [
        {
          "file_id": "1fpiHY_g7b1-bs_sSRWcbiw9qv4eDU4QZ",
          "timestamp": 1628275335747
        },
        {
          "file_id": "https://github.com/tensorflow/text/blob/master/docs/tutorials/transformer.ipynb",
          "timestamp": 1628273726995
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
